#include "kernel_operator.h"
#define OP_INIT_RUN()                           \
    op.Init(x, y, z,                            \
            tiling_data.smallCoreDataNum,       \
            tiling_data.bigCoreDataNum,         \
            tiling_data.finalBigTileNum,        \
            tiling_data.finalSmallTileNum,      \
            tiling_data.tileDataNum,            \
            tiling_data.smallTailDataNum,       \
            tiling_data.bigTailDataNum,         \
            tiling_data.tailBlockNum);          \
    op.Process()
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_X, typename TYPE_Z>
class KernelCast {
    using T = TYPE_X;
public:
    __aicore__ inline KernelCast() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t smallCoreDataNum,
        uint32_t bigCoreDataNum, uint32_t finalBigTileNum,
        uint32_t finalSmallTileNum, uint32_t tileDataNum,
        uint32_t smallTailDataNum, uint32_t bigTailDataNum,
        uint32_t tailBlockNum)
    {
        ASSERT(AscendC::GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalBufferIndex = bigCoreDataNum * AscendC::GetBlockIdx();
        this->tileDataNum = tileDataNum;
        if (coreNum < tailBlockNum) {
          this->coreDataNum = bigCoreDataNum;
          this->tileNum = finalBigTileNum;
          this->tailDataNum = bigTailDataNum;
        }
        else {
          this->coreDataNum = smallCoreDataNum;
          this->tileNum = finalSmallTileNum;
          this->tailDataNum = smallTailDataNum;
          globalBufferIndex -= (bigCoreDataNum - smallCoreDataNum) *
                                (AscendC::GetBlockIdx() - tailBlockNum);
        }

        xGm.SetGlobalBuffer((__gm__ TYPE_X*)x + globalBufferIndex, this->coreDataNum);
        zGm.SetGlobalBuffer((__gm__ TYPE_Z*)z + globalBufferIndex, this->coreDataNum);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_X));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_Z));

        // pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        // pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            if (i == this->tileNum - 1) {
              this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        AscendC::DataCopy(xLocal, xGm[progress * this->tileDataNum], this->processDataNum);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        AscendC::LocalTensor<TYPE_Z> zLocal = outQueueZ.AllocTensor<TYPE_Z>();

        // AscendC::LocalTensor<half> p1Local = tmp1.Get<half>();
        // AscendC::LocalTensor<float> p2Local = tmp2.Get<float>();

        AscendC::Cast(zLocal, xLocal, AscendC::RoundMode::CAST_NONE, this->processDataNum);

        outQueueZ.EnQue<TYPE_Z>(zLocal);
        inQueueX.FreeTensor(xLocal);
        // tmp1.FreeTensor(p1Local);
        // tmp2.FreeTensor(p2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<TYPE_Z> zLocal = outQueueZ.DeQue<TYPE_Z>();
        AscendC::DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN,  BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    // AscendC::TBuf<AscendC::QuePosition::VECCALC> tmp1, tmp2;
    AscendC::GlobalTensor<TYPE_X> xGm;
    AscendC::GlobalTensor<TYPE_Z> zGm;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};

extern "C" __global__ __aicore__ void cast_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                                 GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    
    if (TILING_KEY_IS(1))   { KernelCast<half, float> op;    OP_INIT_RUN(); }   // 0+1
    if (TILING_KEY_IS(2))   { KernelCast<half, uint8_t> op; OP_INIT_RUN(); }   // 0+2
    if (TILING_KEY_IS(3))   { KernelCast<half, int8_t> op;  OP_INIT_RUN(); }   // 0+3
    if (TILING_KEY_IS(4))   { KernelCast<half, int16_t> op; OP_INIT_RUN(); }   // 0+4
    if (TILING_KEY_IS(5))   { KernelCast<half, int32_t> op; OP_INIT_RUN(); }   // 0+5
    // if (TILING_KEY_IS(6))   { KernelCast<half, int64_t> op; OP_INIT_RUN(); }   // 0+6

    if (TILING_KEY_IS(10))  { KernelCast<float, half> op;    OP_INIT_RUN(); }  // 10+0
    // if (TILING_KEY_IS(12))  { KernelCast<float, uint8_t> op; OP_INIT_RUN(); }  // 10+2
    // if (TILING_KEY_IS(13))  { KernelCast<float, int8_t> op;  OP_INIT_RUN(); }  // 10+3
    if (TILING_KEY_IS(14))  { KernelCast<float, int16_t> op; OP_INIT_RUN(); }  // 10+4
    if (TILING_KEY_IS(15))  { KernelCast<float, int32_t> op; OP_INIT_RUN(); }  // 10+5
    if (TILING_KEY_IS(16))  { KernelCast<float, int64_t> op; OP_INIT_RUN(); }  // 10+6

    if (TILING_KEY_IS(20))  { KernelCast<uint8_t, half> op;    OP_INIT_RUN(); } // 20+0
    // if (TILING_KEY_IS(21))  { KernelCast<uint8_t, float> op;   OP_INIT_RUN(); } // 20+1
    // if (TILING_KEY_IS(23))  { KernelCast<uint8_t, int8_t> op;  OP_INIT_RUN(); } // 20+3
    // if (TILING_KEY_IS(24))  { KernelCast<uint8_t, int16_t> op; OP_INIT_RUN(); } // 20+4
    // if (TILING_KEY_IS(25))  { KernelCast<uint8_t, int32_t> op; OP_INIT_RUN(); } // 20+5
    // if (TILING_KEY_IS(26))  { KernelCast<uint8_t, int64_t> op; OP_INIT_RUN(); } // 20+6

    if (TILING_KEY_IS(30))  { KernelCast<int8_t, half> op;    OP_INIT_RUN(); } // 30+0
    // if (TILING_KEY_IS(31))  { KernelCast<int8_t, float> op;   OP_INIT_RUN(); } // 30+1
    // if (TILING_KEY_IS(32))  { KernelCast<int8_t, uint8_t> op; OP_INIT_RUN(); } // 30+2
    // if (TILING_KEY_IS(34))  { KernelCast<int8_t, int16_t> op; OP_INIT_RUN(); } // 30+4
    // if (TILING_KEY_IS(35))  { KernelCast<int8_t, int32_t> op; OP_INIT_RUN(); } // 30+5
    // if (TILING_KEY_IS(36))  { KernelCast<int8_t, int64_t> op; OP_INIT_RUN(); } // 30+6

    if (TILING_KEY_IS(40))  { KernelCast<int16_t, half> op;    OP_INIT_RUN(); } // 40+0
    if (TILING_KEY_IS(41))  { KernelCast<int16_t, float> op;   OP_INIT_RUN(); } // 40+1
    // if (TILING_KEY_IS(42))  { KernelCast<int16_t, uint8_t> op; OP_INIT_RUN(); } // 40+2
    // if (TILING_KEY_IS(43))  { KernelCast<int16_t, int8_t> op;  OP_INIT_RUN(); } // 40+3
    // if (TILING_KEY_IS(45))  { KernelCast<int16_t, int32_t> op; OP_INIT_RUN(); } // 40+5
    // if (TILING_KEY_IS(46))  { KernelCast<int16_t, int64_t> op; OP_INIT_RUN(); } // 40+6

    if (TILING_KEY_IS(50))  { KernelCast<int32_t, half> op;    OP_INIT_RUN(); } // 50+0
    if (TILING_KEY_IS(51))  { KernelCast<int32_t, float> op;   OP_INIT_RUN(); } // 50+1
    // if (TILING_KEY_IS(52))  { KernelCast<int32_t, uint8_t> op; OP_INIT_RUN(); } // 50+2
    // if (TILING_KEY_IS(53))  { KernelCast<int32_t, int8_t> op;  OP_INIT_RUN(); } // 50+3
    if (TILING_KEY_IS(54))  { KernelCast<int32_t, int16_t> op; OP_INIT_RUN(); } // 50+4
    if (TILING_KEY_IS(56))  { KernelCast<int32_t, int64_t> op; OP_INIT_RUN(); } // 50+6

    // if (TILING_KEY_IS(60))  { KernelCast<int64_t, half> op;    OP_INIT_RUN(); } // 60+0
    if (TILING_KEY_IS(61))  { KernelCast<int64_t, float> op;   OP_INIT_RUN(); } // 60+1
    // if (TILING_KEY_IS(62))  { KernelCast<int64_t, uint8_t> op; OP_INIT_RUN(); } // 60+2
    // if (TILING_KEY_IS(63))  { KernelCast<int64_t, int8_t> op;  OP_INIT_RUN(); } // 60+3
    // if (TILING_KEY_IS(64))  { KernelCast<int64_t, int16_t> op; OP_INIT_RUN(); } // 60+4
    if (TILING_KEY_IS(65))  { KernelCast<int64_t, int32_t> op; OP_INIT_RUN(); } // 60+5
    
}